package org.forum.jstun;

import java.net.Inet4Address;
import java.net.Inet6Address;
import java.net.InetAddress;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.HashMap;
import java.util.LinkedHashMap;


public class StunMessage {
	public static final short TYPE_BINDING_REQUEST 			= 0x0001;
	public static final short TYPE_BINDING_RESPONSE 		= 0x0101;
	public static final short TYPE_BINDING_ERROR_RESPONSE 	= 0x0111;
	public static final short TYPE_SHARED_SECRET_REQUEST 	= 0x0002;
	public static final short TYPE_SHARED_SECRET_RESPONSE 	= 0x0102;
	public static final short TYPE_SHARED_SECRET_ERROR 		= 0x0112;

	public static final short HEADER_SIZE 			= 20;
	public static final short TRANSACTION_ID_SIZE 	= 16;

	// message type, 2 bytes
	private short mType = 0;
	// message length(not including header) 2 bytes
	private short mLength = 0;
	// transaction id, 16 bytes
	private byte[] mTransactionId = null;
	// attributes
	private HashMap<Short, StunAttribute> mAttributes = null;

	public StunMessage() {}
	public StunMessage(short type) {
		mType = type;
	}

	public short getType() {
		return mType;
	}

	public void setType(short type) {
		mType = type;
	}

	public short getLength() {
		return mLength;
	}

	public byte[] getTransactionId() {
		return mTransactionId;
	}

	public void setTransactionId(byte[] tid) {
		mTransactionId = tid;
	}

	/**
	 * append attribute into last
	 * @param attribute
	 * @return : if already exist, return false
	 */
	public boolean appendAttribute(StunAttribute attribute) {
		if (mAttributes==null) mAttributes = new LinkedHashMap<Short, StunAttribute>();
		if (mAttributes.containsKey(attribute.getType())) return false;
		mAttributes.put(attribute.getType(), attribute);
		mLength += StunAttribute.HEADER_SIZE + attribute.getLength();
		return true;
	}

	/**
	 * remove attribute by type
	 * @param attrType
	 * @return : if have no this attribute, return false
	 */
	public boolean removeAttribute(short attrType) {
		if (mAttributes==null) return false;
		StunAttribute attr = mAttributes.remove(attrType);
		if (attr!=null) mLength -= 
				StunAttribute.HEADER_SIZE + attr.getLength();
		return attr!=null;
	}

	public StunAttribute getAttribute(short attrType) {
		if (mAttributes==null) return null;
		return mAttributes.get(attrType);
	}

	public Collection<StunAttribute> getAttributes() {
		if (mAttributes==null) return null;
		 return mAttributes.values();
	}

	public void clear() {
		mType = 0;
		mLength = 0;
		mTransactionId = null;
		if (mAttributes!=null) mAttributes.clear();
	}

	public boolean isEmpty() {
		return mType==0 || mTransactionId==null;
	}

	/**
	 * convert all stun message including header and all attributes to byte array
	 * @return
	 */
	public byte[] getBytes() {
		short length = getLength();
		byte[] ret = new byte[HEADER_SIZE+length];
		// message header
		System.arraycopy(StunAttribute.hostToNetwork(mType), 0, ret, 0, 2);
		System.arraycopy(StunAttribute.hostToNetwork(length), 0, ret, 2, 2);
		System.arraycopy(mTransactionId, 0, ret, 4, TRANSACTION_ID_SIZE);
		// attributes
		short size = HEADER_SIZE;
		if (mAttributes==null) return ret;
		for(StunAttribute attr : mAttributes.values()) {
			byte[] attrByte = attr.getBytes();
			System.arraycopy(attrByte, 0, ret, size, attrByte.length);
			size += attrByte.length;
		}
		return ret;
	}

	/**
	 * generate transaction id, size is 16 bytes.
	 * @param hasMagicCookie : only for RFC5389
	 * @return
	 */
	public static byte[] generateId(boolean hasMagicCookie) {
		byte[] tid = new byte[TRANSACTION_ID_SIZE];
		int rand = 0x2112A442; // magic cookie for RFC5389
		while(!hasMagicCookie && rand==0x2112A442)
			rand = (int) (Math.random() * Integer.MAX_VALUE);
		System.arraycopy(StunAttribute.hostToNetwork(rand), 0, tid, 0, 4);
		for(int i = 4; i < TRANSACTION_ID_SIZE; i += 4) {
			rand = (int) (Math.random() * Integer.MAX_VALUE);
			System.arraycopy(StunAttribute.hostToNetwork(rand), 0, tid, i, 4);
		}
		return tid;
	}

	/**
	 * check if transaction id has magic cookie (if support RFC5389)
	 * @param tid : message transaction id byte[16]
	 * @return
	 */
	public static boolean hasMagicCookie(byte[] tid) {
		int magicCookie = StunAttribute.networkToHost(tid, 0, 4);
		return magicCookie == 0x2112A442;
	}

	/**
	 * compare two transaction id
	 * @param tida
	 * @param tidb
	 * @return
	 */
	public static boolean transactionIdEquals(byte[] tida, byte[] tidb) {
		if (tida==null || tida.length!=TRANSACTION_ID_SIZE ||
			tidb==null || tidb.length!=TRANSACTION_ID_SIZE) return false;
		for(int i=0; i<TRANSACTION_ID_SIZE; ++i)
			if (tida[i] != tidb[i]) return false;
		return true;
	}

	/**
	 * generate message integrity with key, automatically align 64 bytes
	 * @param text
	 * @param key
	 * @return
	 */
	public static StunAttribute generateMessageIntegrity(byte[] text, byte[] key) {
		return new MessageIntegrity(StunAttribute.getHmacSHA1(text, key, 
				StunAttribute.MESSAGE_INTEGRITY_ALIGN));
	}

	/**
	 * Only parse Stun Header (20 bytes)
	 * @param data
	 * @return : length is parse result, not attributes size
	 */
	public static StunMessage parseHeader(byte[] data) {
		if (data==null || data.length<HEADER_SIZE) return null;
		int type = StunAttribute.networkToHost(data, 0, 2);
		int len = StunAttribute.networkToHost(data, 2, 2);
		byte[] tid = new byte[16];
		System.arraycopy(data, 4, tid, 0, TRANSACTION_ID_SIZE);
		StunMessage msg = new StunMessage((short)type);
		msg.mLength = (short)len;
		msg.setTransactionId(tid);
		return msg;
	}

	/**
	 * Only parse one Stun Attribute
	 * @param data
	 * @param offset
	 * @return
	 */
	public static StunAttribute parseAttribute(byte[] data, int offset) {
		if (data==null || data.length<offset+4) return null;
		short type = (short) StunAttribute.networkToHost(data, offset, 2);
		short len = (short) StunAttribute.networkToHost(data, offset+2, 2);
		offset += 4;
		if (data.length<offset+len) return null;
		StunAttribute attr = null;
		switch (type) {
			case StunAttribute.TYPE_MAPPED_ADDRESS:
			case StunAttribute.TYPE_RESPONSE_ADDRESS:
			case StunAttribute.TYPE_SOURCE_ADDRESS:
			case StunAttribute.TYPE_CHANGED_ADDRESS:
			case StunAttribute.TYPE_REFLECTED_FROM:
				attr = new StunAddress(type);
				attr.parseValue(data, offset, len);
				break;

			case StunAttribute.TYPE_CHANGE_REQUEST:
				attr = new ChangeRequest(false, false);
				attr.parseValue(data, offset, len);

			case StunAttribute.TYPE_USERNAME:
			case StunAttribute.TYPE_PASSWORD:
				attr = new UsernamePassword(type, null);
				attr.parseValue(data, offset, len);
				break;

			case StunAttribute.TYPE_MESSAGE_INTEGRITY:
				attr = new MessageIntegrity(null);
				attr.parseValue(data, offset, len);
				break;

			case StunAttribute.TYPE_ERROR_CODE:
				attr = new ErrorCode((short)0, null);
				attr.parseValue(data, offset, len);
				break;

			case StunAttribute.TYPE_UNKNOWN_ATTRIBUTES:
				attr = new UnknwonAttributes(null);
				attr.parseValue(data, offset, len);
				break;

			default:
				// unknown attribute, maybe user-define
				byte[] value = new byte[len];
				System.arraycopy(data, offset, value, 0, len);
				attr = new StunAttribute(type, value);
				break;
		}
		return attr;
	}

	public static StunMessage parse(byte[] data, int length) {
		StunMessage msg = parseHeader(data);
		if (msg==null) return null;
		for(int offset=HEADER_SIZE; offset<length;) {
			StunAttribute attr = parseAttribute(data, offset);
			if (attr==null) break;
			msg.appendAttribute(attr);
			offset += StunAttribute.HEADER_SIZE + attr.getLength();
		}
		return msg;
	}

	/** MappedAddress, ResponseAddress, ChangedAddress
	 *  SourceAddress, ReflectedFrom
	 */
	public static class StunAddress extends StunAttribute {
		private short mIpver;
		private int mPort;	// unsigned short (2 bytes)
		private InetAddress mAddress;

		public StunAddress(short type) {
			super(type);
		}

		public StunAddress(short type, short ipver, 
				int port, InetAddress address) {
			super(type, (short)4);
			mIpver = ipver;
			if (mIpver == StunAttribute.IPV6) super.mLength += 16;
			else if (mIpver == StunAttribute.IPV4) super.mLength += 4;
			mPort = port;
			mAddress = address;
		}

		public short getIpver() {
			return mIpver;
		}

		public int getPort() {
			return mPort;
		}

		public InetAddress getAddress() {
			return mAddress;
		}

		@Override
		public byte[] getBytes() {
			short length = getLength();
			byte[] ret = new byte[2+2+length];
			System.arraycopy(StunAttribute.hostToNetwork(getType()), 0, ret, 0, 2);
			System.arraycopy(StunAttribute.hostToNetwork(length), 0, ret, 2, 2);
			// attribute value
			System.arraycopy(StunAttribute.hostToNetwork(mIpver), 0, ret, 4, 2);
			System.arraycopy(StunAttribute.hostToNetwork((short)mPort), 0, ret, 6, 2);
			System.arraycopy(mAddress.getAddress(), 0, ret, 8, length - 4);
			return ret;
		}

		@Override
		public void parseValue(byte[] data, int offset, int length) {
			if (data==null || length<4 || data.length<offset+length) return;
			super.mLength = (short)length;
			mIpver = (short)StunAttribute.networkToHost(data, offset, 2);
			mPort = StunAttribute.networkToHost(data, offset+2, 2);
			try {
				if (mIpver == StunAttribute.IPV4) {
					byte[] addr = new byte[4];
					System.arraycopy(data, offset+4, addr, 0, addr.length);
					mAddress = Inet4Address.getByAddress(addr);
				} else if (mIpver == StunAttribute.IPV6) {
					byte[] addr = new byte[16];
					System.arraycopy(data, offset+4, addr, 0, addr.length);
					mAddress = Inet6Address.getByAddress(addr);
				} else {
					// unknown ip version
				}
			} catch (Exception e){}
		}
	}

	// ChangeRequest
	public static class ChangeRequest extends StunAttribute {
		// value: 00000000 00000000 00000000 00000AB0
		private boolean mChangeIp;		// flag A
		private boolean mChangePort;	// flag B

		public ChangeRequest(Boolean changeIp, Boolean changePort) {
			// const length 4 bytes
			super(StunAddress.TYPE_CHANGE_REQUEST, (short)4);
			mChangeIp = changeIp;
			mChangePort = changePort;
		}

		public boolean isChangeIp() {
			return mChangeIp;
		}

		public boolean isChangePort() {
			return mChangePort;
		}

		@Override
		public byte[] getBytes() {
			byte[] ret = new byte[2+2+4];
			System.arraycopy(StunAttribute.hostToNetwork(getType()), 0, ret, 0, 2);
			System.arraycopy(StunAttribute.hostToNetwork((short)4), 0, ret, 2, 2);
			// attribute value
			Arrays.fill(ret, 4, 4+4, (byte)0);
			ret[ret.length-1] |= mChangeIp ? 4 : 0;
			ret[ret.length-1] |= mChangePort ? 2 : 0;
			return ret;
		}

		@Override
		public void parseValue(byte[] data, int offset, int length) {
			if (data==null || length!=4 || data.length<offset+4) return;
			mChangeIp = (data[offset+4-1] & 4)!=0;
			mChangePort = (data[offset+4-1] & 2)!=0;
		}
	}

	// USERNAME, PASSWORD
	public static class UsernamePassword extends StunAttribute {
		private String mValue = null;

		public UsernamePassword(short type, String value) {
			super(type);
			if (value==null || value.isEmpty()) return;
			byte[] bValue = value.getBytes();
			int length = bValue.length;
			// must be a multiple of 4 (measured in bytes)
			if ((length & 3)!=0) length = (length+4)&(~3);
			super.mLength = (short)length;
			super.mValue = new byte[length];
			System.arraycopy(bValue, 0, super.mValue, 0, bValue.length);
			Arrays.fill(super.mValue, bValue.length, length, (byte)0);
			mValue = value;
		}

		@Override
		public String toString() {
			return mValue;
		}

		@Override
		public void parseValue(byte[] data, int offset, int length) {
			if (data==null || length<1 || data.length<offset+length) return;
			super.mLength = (short)length;
			super.mValue = new byte[length];
			System.arraycopy(data, offset, super.mValue, 0, length);
			mValue = new String(data, offset, length);
		}
	}

	// MessageIntegrity
	public static class MessageIntegrity extends StunAttribute {
		// 20 bytes (160bit) by HMAC-SHA1, input text is:
		// full stun message including header, not including itself,
		// 64 bytes align with padding 0
		public MessageIntegrity(byte[] value) {
			super(StunAttribute.TYPE_MESSAGE_INTEGRITY, value);
		}

		@Override
		public byte[] getBytes() {
			short length = getLength();
			byte[] ret = new byte[2+2+length];
			System.arraycopy(StunAttribute.hostToNetwork(getType()), 0, ret, 0, 2);
			System.arraycopy(StunAttribute.hostToNetwork(length), 0, ret, 2, 2);
			// attribute value
			if (length>0) System.arraycopy(super.mValue, 0, ret, 4, length);
			return ret;
		}

		@Override
		public void parseValue(byte[] data, int offset, int length) {
			if (data==null || length!=20 || data.length<offset+20) return;
			super.mValue = new byte[20];
			System.arraycopy(data, offset, super.mValue, 0, length);
		}
	}

	// ERROR-CODE
	public static class ErrorCode extends StunAttribute {
		// 00000000 00000000 class number
		// class is byte range [1,6], represents the hundreds
		// number is byte range [0, 99], represents code modulo 100.
		private short mCode;
		private String mName;
		public ErrorCode(short code, String name) {
			super(StunAttribute.TYPE_ERROR_CODE);
			mCode = code;
			mName = name;
		}

		@Override
		public short getLength() {
			int length = 4;
			if (mName!=null) length += mName.getBytes().length;
			// must be a multiple of 4 (measured in bytes)
			if ((length & 3)!=0) length = (length+4)&(~3);
			return (short)length;
		}

		@Override
		public byte[] getBytes() {
			short length = getLength();
			byte[] ret = new byte[2+2+length];
			System.arraycopy(StunAttribute.hostToNetwork(getType()), 0, ret, 0, 2);
			System.arraycopy(StunAttribute.hostToNetwork(length), 0, ret, 2, 2);
			// attribute value
			// prefix padding 0 (2 bytes)
			System.arraycopy(StunAttribute.hostToNetwork((short)0), 0, ret, 4, 2);
			// code 2 bytes
			ret[6] = (byte)(mCode/100);
			ret[7] = (byte)(mCode%100);
			// name
			if (length>0) {
				byte[] name = mName.getBytes();
				System.arraycopy(name, 0, ret, 8, name.length);
				// must be a multiple of 4 (measured in bytes)
				Arrays.fill(ret, 8+name.length, 8+length, (byte)0);
			}
			return ret;
		}

		@Override
		public void parseValue(byte[] data, int offset, int length) {
			if (data==null || length<4 || data.length<offset+length) return;
			mCode = (short)(data[offset+2]*100);
			mCode += (short)data[offset+3];
			if (length>4) mName = new String(data, offset+4, length-4);
			else mName = null;
		}
	}

	// UnknwonAttributes
	public static class UnknwonAttributes extends StunAttribute {
		// when error code is 420 EC_UNKNOWN_ATTRIBUTE
		private ArrayList<Short> mAttributes;

		public UnknwonAttributes(ArrayList<Short> attributes) {
			super(StunAttribute.TYPE_UNKNOWN_ATTRIBUTES);
			mAttributes = attributes;
		}

		@Override
		public short getLength() {
			int length = mAttributes!=null ? mAttributes.size()*2 : 0;
			// must be a multiple of 4 (measured in bytes)
			if ((length & 3) != 0) length = (length+4) & (~3);
			return (short) length;
		}

		@Override
		public byte[] getBytes() {
			short length = getLength();
			byte[] ret = new byte[2+2+length];
			System.arraycopy(StunAttribute.hostToNetwork(getType()), 0, ret, 0, 2);
			System.arraycopy(StunAttribute.hostToNetwork(length), 0, ret, 2, 2);
			if (length==0) return ret;
			// attribute value
			short size = 4;
			for(short type:mAttributes) {
				System.arraycopy(StunAttribute.hostToNetwork(type), 0, ret, size, 2);
				size += 2;
			}
			// repeat last attribute if it isn't a multiple of 4 (measured in bytes)
			if (size-4<length) System.arraycopy(ret, size-2, ret, size, 2);
			return ret;
		}

		@Override
		public void parseValue(byte[] data, int offset, int length) {
			if (data==null || length<0 || data.length<offset+length) return;
			if (mAttributes!=null) mAttributes.clear();
			else mAttributes = new ArrayList<Short>();
			for(int pos=offset; pos<offset+length; pos+=4) {
				mAttributes.add((short)StunAttribute.networkToHost(data, pos, 2));
				mAttributes.add((short)StunAttribute.networkToHost(data, pos+2, 2));
			}
		}
	}

}